#include <pinocchio/fwd.hpp>
#include <pinocchio/multibody/model.hpp>
#include <pinocchio/parsers/urdf.hpp>

#include <iostream>
#include <string>
#include <ros/ros.h>
#include <sensor_msgs/JointState.h>
#include <std_msgs/Float64.h>

#include <ocs2_pinocchio_interface/PinocchioInterface.h>
#include <ocs2_oc/oc_problem/OptimalControlProblem.h>
#include <ocs2_core/misc/LoadData.h>
#include <ocs2_core/cost/QuadraticStateInputCost.h>
#include <ocs2_core/cost/QuadraticStateCost.h>
#include <ocs2_core/constraint/LinearStateInputConstraint.h>
#include <ocs2_core/penalties/augmented/SlacknessSquaredHingePenalty.h>
#include <ocs2_core/augmented_lagrangian/AugmentedLagrangian.h>
#include <ocs2_ddp/GaussNewtonDDP_MPC.h>
#include <ocs2_oc/rollout/TimeTriggeredRollout.h>
#include <ocs2_core/initialization/DefaultInitializer.h>
#include <ocs2_mpc/SystemObservation.h>
#include <ocs2_msgs/mpc_target_trajectories.h>
#include <ocs2_ros_interfaces/common/RosMsgConversions.h>
#include <ocs2_ros_interfaces/synchronized_module/RosReferenceManager.h>
#include <ocs2_oc/synchronized_module/ReferenceManager.h>
#include <ocs2_msgs/mpc_observation.h>
#include <ocs2_mpc/MPC_MRT_Interface.h>
#include <ocs2_core/thread_support/ExecuteAndSleep.h>

#include "class03/CartpoleDynamics.h"

using namespace ocs2;
int main(int argc, char **argv)
{
    ros::init(argc, argv, "cartpole_nmpc_controller");
    // 模型
    // 加载urdf
    ros::NodeHandle nh;//ros的代理人，可以处理ros的订阅或者发布，每个文件都可以有一个，不必要在文件之间传递
    std::string urdfFile;
    std::string configFile;
    nh.getParam("urdf_file", urdfFile); // 把“ ”中的 文件路径 加载到后面的变量中
    nh.getParam("config_file", configFile);//必须在.launch文件中写好，不然拿不到；根据执行的那个launch决定，后执行的覆盖先执行的
    //std::cout << "urdf_file: " << urdfFile << std::endl;

    pinocchio::Model model;
    pinocchio::urdf::buildModel(urdfFile, model); //（把urdfFile路径中的urdf文件加载到model中）
    PinocchioInterface pinocchioInterface(model); //（？）
    //std::cout << "model.nq" << model.nq << std::endl;
    //std::cout << "model.nv" << model.nq << std::endl;

    OptimalControlProblem ocp; // 优化控制问题的描述

    // 优化控制问题
    // 动力学
    ocp.dynamicsPtr = std::make_unique<CartpotDynamics>(pinocchioInterface);
    // 损失方程

    int stateDim = 0;//状态维度，用config.info加载
    int inputDim = 0;
    loadData::loadCppDataType(configFile, "StateDim", stateDim);//configFile为前文中在launch里加载的config文件路径
    loadData::loadCppDataType(configFile, "InputDim", inputDim);//把configFile路径中的"InputDim"数值加载到inputDim中
    //std::cout<<"stateDim:"<<stateDim<<std::endl;
    //std::cout<<"inputDim:"<<inputDim<<std::endl;

    //matrix_t是向量/矩阵，是osc2提供的
    matrix_t Q = matrix_t::Identity(stateDim, stateDim);//Indentity 单位矩阵
    matrix_t QF = matrix_t::Identity(stateDim, stateDim);
    matrix_t R = matrix_t::Identity(inputDim, inputDim);
    loadData::loadEigenMatrix(configFile, "Q", Q);
    loadData::loadEigenMatrix(configFile, "QF", QF);
    loadData::loadEigenMatrix(configFile, "R", R);
    //std::cout << "Q:" << Q << std::endl;
    //std::cout << "R:" << R << std::endl;

    ocp.costPtr->add("cartpoleStateInputCost", std::make_unique<QuadraticStateInputCost>(Q, R));//？
    ocp.finalCostPtr->add("cartpoleStateCost", std::make_unique<QuadraticStateCost>(QF));
/*********约束***********/
    scalar_t FMax = 0;
    loadData::loadCppDataType(configFile, "MaxForce", FMax);
    //得到约束，此处根据文档得到C、D、e的取值，getConstraint即为文档中的value
    auto getConstraint = [&]()
    {
        int constraintDim = 2;
        matrix_t C = matrix_t::Zero(constraintDim, stateDim);
        matrix_t D = matrix_t::Zero(constraintDim, inputDim);
        D << 1, -1;
        vector_t e = vector_t::Zero(2);
        e << FMax, FMax;
        return std::unique_ptr<StateInputConstraint>(new LinearStateInputConstraint(e, C, D));
    };
    //得到惩罚，具体原理
    auto getPenalty = [&]()
    {
        scalar_t scale = 1;
        scalar_t stepSize = 1;
        loadData::loadCppDataType(configFile, "MaxForceConstrainPenalty.scale", scale);
        loadData::loadCppDataType(configFile, "MaxForceConstrainPenalty.stepSize", stepSize);
        //std::cout<<"scale:"<<scale<<std::endl;
        //std::cout<<"stepSize:"<<stepSize<<std::endl;

        augmented::SlacknessSquaredHingePenalty::Config config(scale, stepSize);//把惩罚输入
        return augmented::SlacknessSquaredHingePenalty::create(config);
    };
    //把约束和惩罚输入求解器***？
    ocp.inequalityLagrangianPtr->add("inputMaxConstarint", create(getConstraint(), getPenalty()));

    // 创建求解器
    auto rolloutSettings = rollout::loadSettings(configFile, "rollout");
    TimeTriggeredRollout rollout(*ocp.dynamicsPtr, rolloutSettings);
    auto ddpSettings = ddp::loadSettings(configFile, "ddp"); // ddp牛顿迭代法
    auto mpcSettings = mpc::loadSettings(configFile, "mpc");
    DefaultInitializer initializer(inputDim); // 初始化器需要知道输入是几维度
    GaussNewtonDDP_MPC mpc(mpcSettings, ddpSettings, rollout, ocp, initializer);
    MPC_MRT_Interface mpcMRTInterface(mpc);

    vector_t initState = vector_t::Zero(stateDim);
    vector_t initInput = vector_t::Zero(inputDim);
    vector_t finState = vector_t::Zero(stateDim);

    loadData::loadEigenMatrix(configFile, "initState", initState);
    loadData::loadEigenMatrix(configFile, "finalState", finState);
    //std::cout << "initState:" << initState << std::endl;
    //std::cout << "finState:" << finState << std::endl;
    

    TargetTrajectories targetTrajectories{{0, 1},
                                          {initState, finState},
                                          {initInput, initInput}};

    auto referenceManager = std::make_shared<ReferenceManager>(targetTrajectories);
    auto rosReferenceManager = std::make_shared<RosReferenceManager>("cartpole", referenceManager);
    mpc.getSolverPtr()->setReferenceManager(rosReferenceManager);
    rosReferenceManager->subscribe(nh);

    // 给观测状态
    SystemObservation observation;//一个容器，存放着模式、时间、状态和输入
    while (ros::Time::now().toSec() == 0)
        ;
    auto startTime = ros::Time::now();//时间基准，运行时初始化一次

    auto jointStateSub = nh.subscribe<sensor_msgs::JointState>(//话题订阅，每次状态更新运行一次
        "/cartpole/joint_states", 1,
        [&](const sensor_msgs::JointState::ConstPtr &msg)
        {
            observation.time = ros::Time::now().toSec() - startTime.toSec();//此处时间记录着每次状态更新的时间
            observation.state = vector_t::Zero(stateDim); // 初始化观测器阶数
            // 当模型变化时，下面要改变（改变倒立摆阶数）
            observation.state << 
                msg->position[0], // 往观测器装东西
                msg->position[1],
                msg->position[2],
                msg->position[3],
                msg->velocity[0],
                msg->velocity[1],
                msg->velocity[2],
                msg->velocity[3];
            observation.input = vector_t::Zero(inputDim);
        });
    // 给目标

    // 求解优化问题
    auto commandPub = nh.advertise<std_msgs::Float64>(
        "/cartpole/cart_effort_controller/command", 1);
    auto observationPub = nh.advertise<ocs2_msgs::mpc_observation>(
        "/cartpole/mpc_observation",1);




    auto rate = ros::Rate(1000);

    bool mpcRunning = false;
    bool controllrtRunning = true;
    auto mpcThread = std::thread([&](){
        while(controllrtRunning)
        {
            if(mpcRunning)
            {
                executeAndSleep([&](){
                    mpcMRTInterface.advanceMpc();},//其中有mpc_.run
                    mpcSettings.mpcDesiredFrequency_);
            }
            
        }
    });
    while(observation.time == 0)
    {
        rate.sleep();
        ros::spinOnce();//查一下ros服务器
    }
    std::cout<<"first observation received"<<std::endl;
    mpcMRTInterface.setCurrentObservation(observation);//收到状态后传给MRT
    mpcRunning = true;//此时线程激活
    while(!mpcMRTInterface.initialPolicyReceived())//判断初始轨迹是否接受到
    {
        rate.sleep();//等待第一遍算好
    }

    mpcMRTInterface.updatePolicy();
    auto firstPolicy = mpcMRTInterface.getPolicy();//获得第一次轨迹
    for(int i=0;i<firstPolicy.timeTrajectory_.size();i++)
    {
        std::cout<<"time:"<<firstPolicy.timeTrajectory_[i]<<"\t";
        std::cout<<"state:"<<firstPolicy.stateTrajectory_[i].transpose()<<"\t";
        std::cout<<"input:"<<firstPolicy.inputTrajectory_[i].transpose()<<std::endl;
    }

    while (ros::ok())
    {

        ros::spinOnce();
        //更新MRT的obsvation
        mpcMRTInterface.setCurrentObservation(observation);
        //更新policy
        mpcMRTInterface.updatePolicy();
        //获取当前控制指令
        SystemObservation optimizedObservation;
        mpcMRTInterface.evaluatePolicy(observation.time,observation.state,
                                        optimizedObservation.state,
                                        optimizedObservation.input,
                                        optimizedObservation.mode);//从算出的轨迹中处理后，找到最近要跟随的目标点

        
        // 发布控制指令
        std_msgs::Float64 msgs;
        msgs.data = optimizedObservation.input[0];
        commandPub.publish(msgs);
        // 发布当前状态
        auto observation_msgs = ros_msg_conversions::createObservationMsg(observation);
        observationPub.publish(observation_msgs);
        // std::cout<<123<<std::endl;
        rate.sleep();
    }

    // 创建ros节点

    return 0;
}
